Description

This PR adds FlagOS Triton implementations for fused RoPE operators that were previously missing from the FlagOS backend.

The implementation follows the CUDA backend behavior for regular RoPE and fused QKV RoPE paths, including forward and backward execution, multiple QKV layouts,
interleaved and non-interleaved rotary layouts, start position offsets, THD variable-length sequences, and context-parallel position handling.

Fixes: N/A

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Added Triton implementations for:
  • fused_rope_forward
  • fused_rope_backward
  • fused_qkv_rope_forward
  • fused_qkv_rope_backward
• Registered the fused RoPE operators in the FlagOS backend.
• Added backend methods to expose the new FlagOS fused RoPE implementations.
• Added tests for fused RoPE and fused QKV RoPE covering:
  • NVTE_SBHD, NVTE_BSHD, and NVTE_THD
  • interleaved and non-interleaved rotary layouts
  • forward and backward paths
  • start position offsets
  • context-parallel position mapping
  • QKV split handling, including GQA-style splits
• Updated plugin test discovery to include the fused RoPE test suite.
• Improved tests to compare FlagOS outputs against a PyTorch reference and, when available, the CUDA vendor backend.

Validation

• bash ./qa/format.sh
• python3 transformer_engine/plugin/tests/test_fused_rope.py

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes